/*
手搓一个心跳机制的服务器
把心跳机制的整个过程搞透彻了再并到项目中开发
*/

#include <arpa/inet.h>
#include <assert.h>
#include <pthread.h>
#include <sys/epoll.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include <iostream>
#include <mutex>
#include <unordered_map>
#include <vector>
using namespace std;

// 服務器socket初始化

// 全局变量
int sfd = -1;
int efd = -1;
int udpfd = -1;
struct sockaddr_in saddr, udpaddr;
struct epoll_event listenEvent, connectEvent;
struct epoll_event events[1024] = {0};

// vector<struct sockaddr_in> clients;
unordered_map<int, struct sockaddr_in> clients;
unordered_map<int, int> heartBeatCnt;
void socketInit() {
    // scoket
    sfd = socket(AF_INET, SOCK_STREAM, 0);
    // bind
    saddr.sin_family = AF_INET;
    saddr.sin_addr.s_addr = INADDR_ANY;
    saddr.sin_port = htons(9999);
    int ret = bind(sfd, (struct sockaddr *)&saddr, sizeof(saddr));
    assert(!ret);
    // listen
    ret = listen(sfd, 6);
    assert(!ret);
    cout << "socket...bind...listen...done" << endl;
}
void epollInit() {
    efd = epoll_create(1024);
    listenEvent.events |= EPOLLIN;
    listenEvent.data.fd = sfd;
    epoll_ctl(efd, EPOLL_CTL_ADD, sfd, &listenEvent);
}
void udpInit() {
    udpfd = socket(AF_INET, SOCK_DGRAM, 0);
    udpaddr.sin_family = AF_INET;
    udpaddr.sin_addr.s_addr = INADDR_ANY;
    udpaddr.sin_port = htons(9999);
    int ret = bind(udpfd, (struct sockaddr *)&udpaddr, sizeof(udpaddr));
    assert(!ret);
    cout << "udp通信准备就绪" << endl;
}
int handleListen(struct sockaddr_in caddr, socklen_t len) {
    int cfd = accept(sfd, (struct sockaddr *)&caddr, &len);
    assert(cfd);
    connectEvent.events |= EPOLLIN;
    connectEvent.data.fd = cfd;
    epoll_ctl(efd, EPOLL_CTL_ADD, cfd, &connectEvent);
    cout << "分配连接成功" << endl;
    clients.insert({cfd, caddr});
    return cfd;
}

/*
心跳机制需要两个额外的子线程：一个负责接受客户端发送的UDP包，一个负责轮询处理心跳
*/
mutex _mtx;
void *udpLoop(void *) {
    udpInit();
    while (true) {
        // 接收数据
        struct sockaddr_in caddr;
        socklen_t len = 0;
        char buf[128] = {0};
        int l = recvfrom(udpfd, buf, sizeof(buf), 0, (struct sockaddr *)&caddr,
                         &len);
        // if (l < 0) {
        //     // error
        //     cout << "error in udp recvfrom function" << endl;
        //     break;
        // }
        int idx = atoi(buf);
        // 上锁访问
        if (idx > 0) {
            unique_lock<mutex> lk(_mtx);
            heartBeatCnt[idx]--;
            cout << idx << " : "
                 << "heartbeat[" << heartBeatCnt[idx] << "]" << endl;
        }
    }
}
void *heartBeatLoop(void *) {
    while (true) {
        sleep(1);
        {
            unique_lock<mutex> lk(_mtx);
            for (auto iter = heartBeatCnt.begin(); iter != heartBeatCnt.end();
                 ++iter) {
                iter->second++;
                if (iter->second >= 5) {
                    // 调用tcp发送下线通知？
                    cout << iter->first << " offline " << endl;
                    char sendBuf[128] = "offline";
                    send(sfd, sendBuf, sizeof(sendBuf), 0);
                    close(iter->first);
                    iter = heartBeatCnt.erase(iter);
                }
            }
        }
    }
}

int main(int argc, char const *argv[]) {
    //
    socketInit();
    // epoll 的部分
    epollInit();
    // epoll 监听
    // 创建两个子线程
    pthread_t p1, p2;
    pthread_create(&p1, NULL, udpLoop, NULL);
    pthread_create(&p2, NULL, heartBeatLoop, NULL);
    while (true) {
        int sz = epoll_wait(efd, events, 1024, -1);
        for (int i = 0; i < sz; ++i) {
            if (events[i].data.fd == sfd) {
                // 分发连接'
                struct sockaddr_in caddr;
                socklen_t len = 0;
                int cfd = handleListen(caddr, len);
            } else if (events[i].events & EPOLLIN) {
                // 读事件
                char buf[128] = {0};
                int fd = events[i].data.fd;
                int len = read(fd, buf, sizeof(buf));
                if (len == 0) {
                    close(events[i].data.fd);
                    clients.erase(events[i].data.fd);
                    continue;
                }
                cout << "read : " << buf << " | " << fd << endl;
            } else if (events[i].events & EPOLLOUT) {
                // 写事件
            } else if (events[i].events & (EPOLLRDHUP | EPOLLHUP | EPOLLERR)) {
                // 断开连接
                close(events[i].data.fd);
                clients.erase(events[i].data.fd);
            }
        }
        //
    }
    // 回收线程资源
    pthread_join(p1, NULL);
    pthread_join(p2, NULL);
    return 0;
}
